package com.springcloud.gateway.loadBalance;

import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.ClientConfigEnabledRoundRobinRule;
import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.Server;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * @author jeason
 */
@Slf4j
@Component
public class SameIpBalanceRule extends ClientConfigEnabledRoundRobinRule {

    public SameIpBalanceRule(ILoadBalancer lb) {
        this.setLoadBalancer(lb);
    }

    public SameIpBalanceRule() {

    }

    @Override
    public void initWithNiwsConfig(IClientConfig iClientConfig) {

    }

    @Override
    public Server choose(Object ip) {
        log.info("client ip:{}", ip);
        List<Server> servers = this.getLoadBalancer().getReachableServers();
        if (servers.isEmpty()) {
            return null;
        }
        if (servers.size() == 1) {
            return servers.get(0);
        }
        return sameIpChoose(servers, ip);
    }

    private Server sameIpChoose(List<Server> servers, Object ip) {
        for (int i = 0; i < servers.size(); i++) {
            Server server = servers.get(i);
            String host = server.getHost();
            if (StringUtils.equals((CharSequence) ip, host)) {
                return server;
            }
        }
        return super.choose(ip);
    }
}

